package ij3d.image3d;

import ij.*;
import ij.measure.*;
import ij.gui.*;
import ij.process.*;


import java.util.*;
import java.awt.*;

/**
 *  Description of the Class
 *
 *@author     fabrice cordelieres
 *@created    6 mars 2009
 */
public class Segment3D {

    ImagePlus toSegment;
    int lowThreshold;
    int highThreshold;
    double minSizeObject;
    double maxSizeObject;
    String name;

    // FABRICE
    int[] imgArray, objID, IDcount, surfList;
    boolean[] IDisAtEdge;
    boolean[] isSurf;
    int width = 1, height = 1, nbSlices = 1, length = 1, depth = 8;
    int minSize, maxSize, nbObj = 0, nbSurfPix = 0;
    boolean sizeFilter = true, exclude = false, redirect = false, closeImg = false, showMaskedImg = true;
    boolean foundObjects = false;
    int thr = 0;
    Calibration cal;

    /**
     *  Constructor for the Segment3D object
     *
     *@param  img      The image to segment
     *@param  lowthr   Low threshold
     *@param  highthr  High Threshold
     */
    public Segment3D(ImagePlus img, int lowthr, int highthr) {
        this.width = img.getWidth();
        this.height = img.getHeight();
        this.nbSlices = img.getNSlices();
        this.length = this.width * this.height * this.nbSlices;
        this.cal = img.getCalibration();
        this.thr = lowthr;
        this.minSize = 1;
        this.maxSize = length;
        this.sizeFilter = true;
        this.exclude = false;
        this.redirect = false;
        this.name = img.getShortTitle();

        // TB
        this.lowThreshold = lowthr;
        this.highThreshold = highthr;

        if (depth != 8 && depth != 16) {
            throw new IllegalArgumentException("Counter3D class expects 8- or 16-bits images only");
        }

        this.nbObj = this.length;

        this.imgArray = new int[this.length];

        this.imgArrayModifier(img);
    }

    /**
     *  Generates the connexity analysis.
     */
    public void findObjects() {
        //First ID attribution
        int currID = 0;
        int currPos = 0;
        int minID = 0;
        int surfPix = 0;
        int neigbNb = 0;
        int pos;
        int currPixID;
        int neigbX;
        int neigbY;
        int neigbZ;

        long start = System.currentTimeMillis();
        /*
         *  Finding the structures:
         *  The minID tag is initialized with the current value of tag (currID).If thresholded,
         *  the neighborhood of the current pixel is collected. For each of those 13 pixels,
         *  the value is retrieved and tested against minID: only the minimum of the two is kept.
         *  As anterior pixels have already been tagged, only two possibilities may exists:
         *  1-The minimum is currID: we start a new structure and currID should be incremented
         *  2-The minimum is not currID: we continue an already existing structure
         *  Each time a new pixel is tagged, a counter of pixels in the current tag is incremented.
         */
        this.objID = new int[this.length];

        for (int z = 1; z <= this.nbSlices; z++) {
            for (int y = 0; y < this.height; y++) {
                for (int x = 0; x < this.width; x++) {
                    if (minID == currID) {
                        currID++;
                    }
                    if (this.imgArray[currPos] != 0) {
                        minID = currID;
                        minID = minAntTag(minID, x, y, z);
                        this.objID[currPos] = minID;
                    }
                    currPos++;
                }
            }
            //IJ.showStatus("Finding structures "+z*100/this.nbSlices+"%");
            IJ.showStatus("Step 1/3: Finding structures");
            IJ.showProgress(z, this.nbSlices);
        }
        IJ.showStatus("");

        this.IDcount = new int[currID];
        for (int i = 0; i < this.length; i++) {
            this.IDcount[this.objID[i]]++;
        }

        this.IDisAtEdge = new boolean[currID];
        Arrays.fill(this.IDisAtEdge, false);
        /*
         *  Connecting structures:
         *  The first tagging of structure may have led to shearing apart pieces of a same structure
         *  This part will connect them back by attributing the minimal retrieved tag among the 13 neighboring
         *  pixels located prior to the current pixel + the centre pixel and will replace all the values of those pixels
         *  by the minimum value.
         */
        this.isSurf = new boolean[this.length];
        currPos = 0;
        minID = 1;

        for (int z = 1; z <= this.nbSlices; z++) {
            for (int y = 0; y < this.height; y++) {
                for (int x = 0; x < this.width; x++) {
                    if (this.imgArray[currPos] != 0) {
                        minID = this.objID[currPos];
                        surfPix = 0;
                        neigbNb = 0;
                        //Find the minimum tag in the neighbours pixels
                        for (neigbZ = z - 1; neigbZ <= z + 1; neigbZ++) {
                            for (neigbY = y - 1; neigbY <= y + 1; neigbY++) {
                                for (neigbX = x - 1; neigbX <= x + 1; neigbX++) {
                                    //Following line is important otherwise objects might be linked from one side of the stack to the other !!!
                                    if (neigbX >= 0 && neigbX < this.width && neigbY >= 0 && neigbY < this.height && neigbZ >= 1 && neigbZ <= this.nbSlices) {
                                        pos = offset(neigbX, neigbY, neigbZ);
                                        if (this.imgArray[pos] != 0) {
                                            if ((this.nbSlices > 1 && ((neigbX == x && neigbY == y && neigbZ == z - 1) || (neigbX == x && neigbY == y && neigbZ == z + 1))) || (neigbX == x && neigbY == y - 1 && neigbZ == z) || (neigbX == x && neigbY == y + 1 && neigbZ == z) || (neigbX == x - 1 && neigbY == y && neigbZ == z) || (neigbX == x + 1 && neigbY == y && neigbZ == z)) {
                                                surfPix++;
                                            }
                                            minID = Math.min(minID, this.objID[pos]);
                                        }
                                        neigbNb++;
                                    }
                                }
                            }
                        }
                        if ((surfPix != 6 && this.nbSlices > 1) || (surfPix != 4 && this.nbSlices == 1)) {
                            this.isSurf[currPos] = true;
                            this.nbSurfPix++;
                        } else {
                            this.isSurf[currPos] = false;
                        }
                        //Replacing tag by the minimum tag found
                        for (neigbZ = z - 1; neigbZ <= z + 1; neigbZ++) {
                            for (neigbY = y - 1; neigbY <= y + 1; neigbY++) {
                                for (neigbX = x - 1; neigbX <= x + 1; neigbX++) {
                                    //Following line is important otherwise objects might be linked from one side of the stack to the other !!!
                                    if (neigbX >= 0 && neigbX < this.width && neigbY >= 0 && neigbY < this.height && neigbZ >= 1 && neigbZ <= this.nbSlices) {
                                        pos = offset(neigbX, neigbY, neigbZ);
                                        if (this.imgArray[pos] != 0) {
                                            currPixID = this.objID[pos];
                                            if (currPixID > minID) {
                                                replaceID(currPixID, minID);
                                            }
                                        }
                                    }
                                }
                            }
                        }

                        //Check if the current particle is touching an edge
                        if (x == 0 || y == 0 || x == this.width - 1 || y == this.height - 1 || (this.nbSlices != 1 && (z == 1 || z == this.nbSlices))) {
                            this.IDisAtEdge[minID] = true;
                        }
                    }
                    currPos++;
                }
            }
            IJ.showStatus("Step 2/3: Connecting structures");
            IJ.showProgress(z, this.nbSlices);
        }
        IJ.showStatus("");

        int newCurrID = 0;

        //Renumbering of all the found objects and update of their respective number of pixels while filtering based on the number of pixels
        for (int i = 1; i < this.IDcount.length; i++) {
            if ((this.IDcount[i] != 0 && this.IDcount[i] >= this.minSize && this.IDcount[i] <= this.maxSize) && (!this.exclude || !(this.exclude && this.IDisAtEdge[i]))) {
                newCurrID++;
                int nbPix = this.IDcount[i];
                replaceID(i, newCurrID);
                this.IDcount[newCurrID] = nbPix;
            } else {
                replaceID(i, 0);
            }
            IJ.showStatus("Step 3/3: Renumbering structures");
            IJ.showProgress(i, currID);
        }
        IJ.showStatus("3D segmentation : " + (System.currentTimeMillis() - start) / 1000.0 + " ms");

        //if (this.redirect) {
        //	this.prepareImgArrayForRedirect();
        //	}
        //if (this.showMaskedImg) {
        //this.buildImg(this.imgArray, null, "Masked image", false, false, false, 0, 0).show();
        this.buildImg(this.objID, null, name + "-3Dseg", false, false, false, 0, 0).show();
        // compute surf
        this.surfList=new int[this.length];
        for (int i=0; i<this.length; i++) this.surfList[i]=this.isSurf[i]?this.objID[i]:0;
        this.buildImg(this.surfList, null, name+"-3Dsurf", false, false, false, 0, 0).show();
        //}

        this.nbObj = newCurrID;
        this.foundObjects = true;

    //this.getObjects();
    }

    /**
     *  Returns the minimum anterior tag among the 13 previous pixels (4 pixels in
     *  2D).
     *
     *@param  x             coordinate x of the current pixel.
     *@param  y             coordinate y of the current pixel.
     *@param  z             coordinate z of the current pixel.
     *@param  initialValue  Description of the Parameter
     *@return               the minimum found anterior tag as an integer.
     */
    private int minAntTag(int initialValue, int x, int y, int z) {
        int min = initialValue;
        int currPos;

        for (int neigbY = y - 1; neigbY <= y + 1; neigbY++) {
            for (int neigbX = x - 1; neigbX <= x + 1; neigbX++) {
                //Following line is important otherwise objects might be linked from one side of the stack to the other !!!
                if (neigbX >= 0 && neigbX < this.width && neigbY >= 0 && neigbY < this.height && z - 1 >= 1 && z - 1 <= this.nbSlices) {
                    currPos = offset(neigbX, neigbY, z - 1);
                    if (this.imgArray[currPos] != 0) {
                        min = Math.min(min, this.objID[currPos]);
                    }
                }
            }
        }

        for (int neigbX = x - 1; neigbX <= x + 1; neigbX++) {
            //Following line is important otherwise objects might be linked from one side of the stack to the other !!!
            if (neigbX >= 0 && neigbX < this.width && y - 1 >= 0 && y - 1 < this.height && z >= 1 && z <= this.nbSlices) {
                currPos = offset(neigbX, y - 1, z);
                if (this.imgArray[currPos] != 0) {
                    min = Math.min(min, this.objID[currPos]);
                }
            }
        }

        //Following line is important otherwise objects might be linked from one side of the stack to the other !!!
        if (x - 1 >= 0 && x - 1 < this.width && y >= 0 && y < this.height && z >= 1 && z <= this.nbSlices) {
            currPos = offset(x - 1, y, z);
            if (this.imgArray[currPos] != 0 && x >= 1 && y >= 0 && z >= 1) {
                min = Math.min(min, this.objID[currPos]);
            }
        }

        return min;
    }

    /**
     *  Returns the index where to find the informations corresponding to pixel (x,
     *  y, z).
     *
     *@param  m  coordinate x of the current pixel
     *@param  n  coordinate y of the current pixel
     *@param  o  coordinate z of the current pixel
     *@return    the index where to find the informations corresponding to pixel
     *      (x, y, z).
     */
    private int offset(int m, int n, int o) {
        if (m + n * this.width + (o - 1) * this.width * this.height >= this.width * this.height * this.nbSlices) {
            return this.width * this.height * this.nbSlices - 1;
        } else {
            if (m + n * this.width + (o - 1) * this.width * this.height < 0) {
                return 0;
            } else {
                return m + n * this.width + (o - 1) * this.width * this.height;
            }
        }
    }

    /**
     *  Replaces one object ID by another within the objID array.
     *
     *@param  oldVal  old value
     *@param  newVal  new value
     */
    private void replaceID(int oldVal, int newVal) {
        if (oldVal != newVal) {
            int nbFoundPix = 0;
            for (int i = 0; i < this.objID.length; i++) {
                if (this.objID[i] == oldVal) {
                    this.objID[i] = newVal;
                    nbFoundPix++;
                }
                if (nbFoundPix == this.IDcount[oldVal]) {
                    i = this.objID.length;
                }
            }
            this.IDcount[oldVal] = 0;
            this.IDcount[newVal] += nbFoundPix;
        }
    }

    /**
     *  Set to zero pixels below the threshold in the "imgArray" arrays.
     *
     *@param  img  The image
     */
    private void imgArrayModifier(ImagePlus img) {
        int index = 0;
        for (int i = 1; i <= this.nbSlices; i++) {
            img.setSlice(i);
            for (int j = 0; j < this.height; j++) {
                for (int k = 0; k < this.width; k++) {
                    this.imgArray[index] = img.getProcessor().getPixel(k, j);
                    if ((this.imgArray[index] < this.lowThreshold) || (this.imgArray[index] > this.highThreshold)) {
                        this.imgArray[index] = 0;
                        this.nbObj--;
                    }
                    index++;
                }
            }
        }
        if (this.closeImg) {
            img.close();
        }
        if (this.nbObj <= 0) {
            IJ.log("No object found in "+name);
            return;
        }
    }

    /**
     *  Generates the ImagePlus based on Counter3D object width, height and number
     *  of slices, the input array and title.
     *
     *@param  imgArray  containing the pixels intensities (integer array).
     *@param  cenArray  containing the coordinates of pixels where the labels
     *      should be put (integer array).
     *@param  title     to attribute to the ImagePlus (string).
     *@param  drawDots  should be true if dots should be drawn instead of a single
     *      pixel for each coordinate of imgArray (boolean).
     *@param  drawNb    should be true if numbers have to be drawn at each
     *      coordinate stored in cenArray (boolean).
     *@param  whiteNb   should be true if numbers have to appear white (boolean).
     *@param  dotSize   size of the dots to be drawn (integer).
     *@param  fontSize  font size of the numbers to be shown (integer).
     *@return           Description of the Return Value
     */
    private ImagePlus buildImg(int[] imgArray, int[] cenArray, String title, boolean drawDots, boolean drawNb, boolean whiteNb, int dotSize, int fontSize) {
        int index = 0;
        int imgDepth = 16;
        float min = imgArray[0];
        float max = imgArray[0];

        for (int i = 0; i < imgArray.length; i++) {
            int currVal = imgArray[i];
            min = Math.min(min, currVal);
            max = Math.max(max, currVal);
        }

        if (max < 256) {
            imgDepth = 8;
        }
        ImagePlus img = NewImage.createImage(title, this.width, this.height, this.nbSlices, imgDepth, 1);

        for (int z = 1; z <= this.nbSlices; z++) {
            IJ.showStatus("Creating the image...");
            img.setSlice(z);
            ImageProcessor ip = img.getProcessor();
            for (int y = 0; y < this.height; y++) {
                for (int x = 0; x < this.width; x++) {
                    int currVal = imgArray[index];
                    if (currVal != 0) {
                        ip.setValue(currVal);
                        if (drawDots) {
                            ip.setLineWidth(dotSize);
                            ip.drawDot(x, y);
                        } else {
                            ip.putPixel(x, y, currVal);
                        }
                    }
                    index++;
                }
            }
        }
        IJ.showStatus("");

        index = 0;
        if (drawNb && cenArray != null) {
            for (int z = 1; z <= this.nbSlices; z++) {
                IJ.showStatus("Numbering objects...");
                img.setSlice(z);
                ImageProcessor ip = img.getProcessor();
                ip.setValue(Math.pow(2, imgDepth));
                ip.setFont(new Font("Arial", Font.PLAIN, fontSize));
                for (int y = 0; y < this.height; y++) {
                    for (int x = 0; x < this.width; x++) {
                        int currVal = cenArray[index];
                        if (currVal != 0) {
                            if (!whiteNb) {
                                ip.setValue(currVal);
                            }
                            ip.drawString("" + currVal, x, y);
                        }
                        index++;
                    }
                }
            }
        }
        IJ.showStatus("");

        img.setCalibration(this.cal);
        img.setDisplayRange(min, max);
        return img;
    }
}

